package main

import (
	"fmt"
	"net"
	"regexp"
	"strings"
	"sync"
)

const (
	serverAddress = "192.168.37.20"
	badServer     = "192.168.102.139"
	serverPort    = 554
)

func getaddress(response []byte) string {
	pattern := `rtsp://([^:]+)`
	re := regexp.MustCompile(pattern)
	matches := re.FindStringSubmatch(string(response))
	if len(matches) > 1 {
		return matches[1]
	}
	return "0"
}

func replaceip(response []byte, ipv4Address string) []byte {
	domainPattern := `rtsp://\b[a-zA-Z0-9_.-]+\.[a-zA-Z]{2,}\b`
	ipPattern := `rtsp://\d+\.\d+\.\d+\.\d+`
	reDomain := regexp.MustCompile(domainPattern)
	reIP := regexp.MustCompile(ipPattern)
	tosend := reDomain.ReplaceAllString(string(response), "rtsp://"+ipv4Address)
	tosend = reIP.ReplaceAllString(tosend, "rtsp://"+ipv4Address)
	return []byte(tosend)
}

func handleStoC(clientSocket, serverSocket net.Conn, wg *sync.WaitGroup) {
	defer wg.Done()
	fmt.Println("开始转发数据")
	buffer := make([]byte, 4096)
	for {
		n, err := serverSocket.Read(buffer)
		if err != nil {
			fmt.Printf("[!]收发错误！关闭连接: %v\n", err)
			serverSocket.Close()
			clientSocket.Close()
			return
		}
		clientSocket.Write(buffer[:n])
	}
}

func handleCtoS(clientSocket, serverSocket net.Conn, wg *sync.WaitGroup) {
	defer wg.Done()
	fmt.Println("开始转发数据")
	buffer := make([]byte, 4096)
	for {
		n, err := clientSocket.Read(buffer)
		if err != nil {
			fmt.Printf("[!]收发错误！关闭连接: %v\n", err)
			serverSocket.Close()
			clientSocket.Close()
			return
		}
		serverSocket.Write(buffer[:n])
	}
}

func describe(firstRequest []byte) []byte {
	newIPPort := fmt.Sprintf("%s:%d", serverAddress, serverPort)
	originalString := string(firstRequest)
	originalString = regexp.MustCompile(`rtsp://\b[a-zA-Z0-9_.-]+\.[a-zA-Z]{2,}\b:\d+`).ReplaceAllString(originalString, "rtsp://"+newIPPort)
	originalString = regexp.MustCompile(`rtsp://\d+\.\d+\.\d+\.\d+:\d+`).ReplaceAllString(originalString, "rtsp://"+newIPPort)
	lines := strings.Split(originalString, "\n")
	lines[0] = strings.Replace(lines[0], "OPTIONS", "DESCRIBE", 1)
	lines[1] = strings.Replace(lines[1], "CSeq: 1", "CSeq: 2", 1)
	lines = append([]string{lines[0], "Accept: application/sdp"}, lines[1:]...)
	return []byte(strings.Join(lines, "\n"))
}

func getnewserver(firstRequest []byte, serverSocket net.Conn) (net.Conn, []byte, string) {
	twoToServer := describe(firstRequest)
	fmt.Printf("继续发送服务器的二次请求:\n%s\n", twoToServer)
	serverSocket.Write(twoToServer)
	buffer := make([]byte, 4096)
	n, _ := serverSocket.Read(buffer)
	response := buffer[:n]
	fmt.Printf("接收到的二次响应:\n%s\n", response)
	targetAddress := getaddress(response)
	if targetAddress == badServer {
		fmt.Println("空的目标服务器关闭")
		return nil, nil, ""
	}
	fmt.Printf("收到转跳命令后，用首次请求数据:\n%s\n", firstRequest)
	fmt.Printf("获取二次响应中的实际目标地址:\n%s\n", targetAddress)
	oneToServer := replaceip(firstRequest, targetAddress)
	serverSocket, err := net.Dial("tcp", fmt.Sprintf("%s:%d", targetAddress, serverPort))
	if err != nil {
		fmt.Printf("目标服务器无响应，关闭\n")
		return nil, nil, ""
	}
	fmt.Printf("重新连接发往目标服务器的请求为:\n%s\n", oneToServer)
	serverSocket.Write(oneToServer)
	n, _ = serverSocket.Read(buffer)
	response = buffer[:n]
	fmt.Printf("重新连接接收到新服务器响应为:\n%s\n", response)
	return serverSocket, response, targetAddress
}

func nonsession(response []byte) bool {
	re := regexp.MustCompile(`Session:\s*(\d+)`)
	matches := re.FindStringSubmatch(string(response))
	return len(matches) == 0
}

func handleEntrance(clientSocket net.Conn) {
	defer clientSocket.Close()
	targetAddress := serverAddress
	serverSocket, err := net.Dial("tcp", fmt.Sprintf("%s:%d", targetAddress, serverPort))
	if err != nil {
		fmt.Println("连接失败")
		return
	}
	defer serverSocket.Close()
	buffer := make([]byte, 1024)
	for i := 0; i < 4; i++ {
		n, err := clientSocket.Read(buffer)
		if err != nil {
			fmt.Println("连接失败")
			return
		}
		request := buffer[:n]
		fmt.Printf("接收到的客户端的请求为:\n%s\n", request)
		clientAddress := getaddress(request)
		toServer := replaceip(request, targetAddress)
		fmt.Printf("发往目标服务器的请求为:\n%s\n", toServer)
		serverSocket.Write(toServer)
		n, _ = serverSocket.Read(buffer)
		response := buffer[:n]
		fmt.Printf("接收到服务器响应为:\n%s\n", response)
		if nonsession(response) {
			fmt.Println("该响应不包含会话ID，需要重新连接")
			serverSocket, response, targetAddress = getnewserver(toServer, serverSocket)
			fmt.Println("新目标服务器已经重新连接")
		}
		toClient := replaceip(response, clientAddress)
		fmt.Printf("发往客户端的响应为:\n%s\n", toClient)
		clientSocket.Write(toClient)
	}
	var wg sync.WaitGroup
	wg.Add(2)
	go handleStoC(clientSocket, serverSocket, &wg)
	go handleCtoS(clientSocket, serverSocket, &wg)
	wg.Wait()
}

func startProxy() {
	proxySocket, err := net.Listen("tcp", fmt.Sprintf(":%d", serverPort))
	if err != nil {
		fmt.Println("监听失败")
		return
	}
	defer proxySocket.Close()
	fmt.Println("开始监听554端口")
	for {
		clientSocket, addr := proxySocket.Accept()
		if err != nil {
			fmt.Println("接受连接失败")
			continue
		}
		fmt.Printf("连接 %v\n", addr)
		go handleEntrance(clientSocket)
	}
}

func main() {
	startProxy()
}
